//
// Created by Administrator on 2022/10/31.
//
#include <onix/memory.h>
#include <onix/assert.h>
#include <onix/debug.h>
#include <onix/string.h>
#include <onix/bitmap.h>
#include <onix/task.h>

typedef struct ards_t
{
    u64 base; // 内存基地址
    u64 size; // 内存长度
    u32 type; // 类型
} _packed ards_t;

#define ZONE_VALID 1    // ards 可用内存区域
#define ZONE_RESERVED 2 // ards 不可用区域

// 获取addr的索引
#define IDX(addr) ((u32)(addr) >> 12)
// 获取页索引idx对应的页开始位置
#define PAGE(idx)((u32)(idx) << 12)
// 获取 addr 的页目录索引
#define DIDX(addr) (((u32)(addr) >> 22) & 0x3FF)
// 获取 addr 的页表索引
#define TIDX(addr) (((u32)(addr) >> 12) & 0x3FF)

#define PDE_MASK 0xFFC00000

#define ASSERT_PAGE(addr) assert(((addr) & 0xFFF) == 0)

// 可用内存基地址
static u32 memory_base = 0;
// 可用内存大小
static u32 memory_size = 0;
// 所有内存页数
static u32 total_pages = 0;
static u32 free_pages = 0;  // 空闲内存页数
// 已使用内存页数
#define used_pages (total_pages - free_pages)

bitmap_t kernel_map;
#define KERNEL_MAP_BITS 0x6000

static void enable_page()
{
    __asm__ volatile("movl %cr0, %eax\n"
                     "orl $0x80000000, %eax\n"
                     "movl %eax, %cr0\n");
}

static void entry_init(page_entry_t *entry, u32 index)
{
    *(u32*)entry = 0;
    entry->present = 1;
    entry->write = 1;
    entry->user = 1;
    entry->index = index;
}

u32 get_cr2()
{
    __asm__ volatile("movl %cr2, %eax\n");
}

u32 get_cr3()
{
    __asm__ volatile("movl %cr3, %eax\n");
}

void set_cr3(u32 pde)
{
    ASSERT_PAGE(pde);
    __asm__ volatile("movl %%eax, %%cr3\n" ::"a"(pde));
}

void memory_init(u32 magic, u32 addr)
{
    u32 count = 0;

    // 如果是 onix loader 进入的内核
    if (magic == ONIX_MAGIC)
    {
        count = *(u32 *)addr;
        ards_t *ptr = (ards_t *)(addr + 4);

        for (size_t i = 0; i < count; i++, ptr++)
        {
            DEBUGK("Memory base 0x%p size 0x%p type %d\n",
                 (u32)ptr->base, (u32)ptr->size, (u32)ptr->type);
            if (ptr->type == ZONE_VALID && ptr->size > memory_size)
            {
                memory_base = (u32)ptr->base;
                memory_size = (u32)ptr->size;
            }
        }
    }

    DEBUGK("ARDS count %d\n", count);
    DEBUGK("Memory base 0x%p\n", (u32)memory_base);
    DEBUGK("Memory size 0x%p\n", (u32)memory_size);

    // memory_size >> 12 可得到该内存区域可分为多少页（每页4K）
    total_pages = IDX(memory_size) + IDX(MEMORY_BASE);
    free_pages = IDX(memory_size);

    DEBUGK("Total pages %d\n", total_pages);
    DEBUGK("Free pages %d\n", free_pages);

    if (memory_size < KERNEL_MEMORY_SIZE)
    {
        panic("System memory is %dM to small, at least %dM needed\n",
              memory_size / MEMORY_BASE, KERNEL_MEMORY_SIZE / MEMORY_BASE);
    }
}

// 可分配物理内存起始位置
static u32 start_page = 0;
// 物理内存数组
static u8 *memory_map;
// 物理内存数组占用的页数
static u32 memory_map_pages;

void memory_map_init()
{
    // 初始化物理内存数组
    memory_map = (u8 *)memory_base;

    // 计算物理内存数组占用的页数
    memory_map_pages = div_round_up(total_pages, PAGE_SIZE);
    DEBUGK("Memory map page count %d\n", memory_map_pages);

    free_pages -= memory_map_pages;

    // 清空物理内存数组
    memset((void *)memory_map, 0, memory_map_pages * PAGE_SIZE);

    // 前 1M 的内存位置 以及 物理内存数组已占用的页，已被占用
    start_page = IDX(MEMORY_BASE) + memory_map_pages;
    for (size_t i = 0; i < start_page; i++)
    {
        memory_map[i] = 1;
    }
    DEBUGK("Total pages %d free pages %d\n", total_pages, free_pages);

    u32 length = (IDX(KERNEL_MEMORY_SIZE) - IDX(MEMORY_BASE)) / 8;
    bitmap_init(&kernel_map, (char*)KERNEL_MAP_BITS, length, IDX(MEMORY_BASE));
    bitmap_scan(&kernel_map, memory_map_pages);
}

void mapping_init()
{
    // 页目录
    page_entry_t *pde = (page_entry_t*)KERNEL_PAGE_DIR;
    memset(pde, 0, PAGE_SIZE);

    idx_t index = 0;
    // 初始化所有内核页表
    for (idx_t didx = 0; didx < (sizeof(KERNEL_PAGE_TABLE) / 4); didx++)
    {
        // 得到页表项地址
        page_entry_t *pte = (page_entry_t*)KERNEL_PAGE_TABLE[didx];
        memset(pte, 0, PAGE_SIZE);

        // 在页目录中设置每个页表的索引地址
        page_entry_t *dentry = &pde[didx];
        entry_init(dentry, IDX((u32)pte));

        // 初始化物理页内存
        for (size_t tidx = 0; tidx < 1024; index++, tidx++) {
            // 第0页不映射，伪造空指针
            if (index == 0) {
                continue;
            }
            page_entry_t *tentry = &pte[tidx];
            entry_init(tentry, index);
            memory_map[index] = 1;
        }
    }

    // 将最后一个页表指向页目录自身，方便修改
    page_entry_t *entry = &pde[1023];
    entry_init(entry, IDX(KERNEL_PAGE_DIR));
    set_cr3((u32)pde);
    enable_page();
}

// 分配一页物理内存
static u32 get_page()
{
    for (size_t i = start_page; i < total_pages; i++)
    {
        // 如果物理内存没有占用
        if (!memory_map[i])
        {
            memory_map[i] = 1;
            free_pages--;
            u32 page = PAGE(i);
            DEBUGK("GET page 0x%p\n", page);
            return page;
        }
    }
    return 0;
}

// 释放一页物理内存
static void put_page(u32 addr)
{
    u32 idx = IDX(addr);
    // 物理引用减一
    memory_map[idx]--;
    // 若为 0，则空闲页加一
    if (!memory_map[idx])
    {
        free_pages++;
    }
    DEBUGK("PUT page 0x%p\n", addr);
}

void memory_test()
{
    u32 *pages = (u32*)(0x200000);
    u32 count = 0x6FE;
    for (size_t i = 0; i < count; ++i) {
        pages[i] = alloc_kpage(1);
        DEBUGK("%#x\n", i);
    }

    for (size_t i = 0; i < count; ++i) {
        free_kpage(pages[i], 1);
    }
}

static page_entry_t *get_pde()
{
    return (page_entry_t*)(0xFFFFF000);
}

static page_entry_t *get_pte(u32 vaddr, bool create)
{
    page_entry_t *pde = get_pde();
    u32 idx = DIDX(vaddr);
    page_entry_t *entry = &pde[idx];

    assert(create || (!create && entry->present));

    page_entry_t *table = (page_entry_t *)(PDE_MASK | (idx << 12));

    if (!entry->present)
    {
        DEBUGK("Get and create page table entry for 0x%p\n", vaddr);
        u32 page = get_page();
        entry_init(entry, IDX(page));
        memset(table, 0, PAGE_SIZE);
    }

    return table;
}

// 刷新虚拟地址 vaddr 的 块表 TLB
static void flush_tlb(u32 vaddr)
{
    __asm__ volatile("invlpg (%0)" ::"r"(vaddr)
    : "memory");
}

// 从位图中扫描 count 个连续的页
static u32 scan_page(bitmap_t *map, u32 count)
{
    int32 index = bitmap_scan(map, count);
    if (index == EOF) {
        panic("Scan page fail!!!");
    }

    u32 addr = PAGE(index);
    DEBUGK("Scan page 0x%p count %d\n", addr, count);
    return addr;
}

// 与 scan_page 相对，重置相应的页
static void reset_page(bitmap_t *map, u32 addr, u32 count)
{
    u32 index = IDX(addr);

    for (size_t i = 0; i < count; i++)
    {
        bitmap_set(map, index + i, 0);
    }
}

// 分配 count 个连续的内核页
u32 alloc_kpage(u32 count)
{
    u32 vaddr = scan_page(&kernel_map, count);
    DEBUGK("ALLOC kernel pages 0x%p count %d\n", vaddr, count);
    return vaddr;
}

// 释放 count 个连续的内核页
void free_kpage(u32 vaddr, u32 count)
{
    reset_page(&kernel_map, vaddr, count);
    DEBUGK("FREE  kernel pages 0x%p count %d\n", vaddr, count);
}

// 将 vaddr 映射物理内存
void link_page(u32 vaddr)
{
    ASSERT_PAGE(vaddr);

    page_entry_t *pte = get_pte(vaddr, TRUE);
    page_entry_t *entry = &pte[TIDX(vaddr)];

    task_t *task = running_task();
    bitmap_t *map = task->vmap;
    u32 index = IDX(vaddr);

    // 如果页面已存在，则直接返回
    if (entry->present)
    {
        assert(bitmap_test(map, index));
        return;
    }

    assert(!bitmap_test(map, index));
    bitmap_set(map, index, TRUE);

    u32 paddr = get_page();
    entry_init(entry, IDX(paddr));
    flush_tlb(vaddr);

    DEBUGK("LINK from 0x%p to 0x%p\n", vaddr, paddr);
}

// 去掉 vaddr 对应的物理内存映射
void unlink_page(u32 vaddr)
{
    ASSERT_PAGE(vaddr);

    page_entry_t *pte = get_pte(vaddr, TRUE);
    page_entry_t *entry = &pte[TIDX(vaddr)];

    task_t *task = running_task();
    bitmap_t *map = task->vmap;
    u32 index = IDX(vaddr);

    if (!entry->present)
    {
        assert(!bitmap_test(map, index));
        return;
    }

    assert(entry->present && bitmap_test(map, index));

    entry->present = FALSE;
    bitmap_set(map, index, FALSE);

    u32 paddr = PAGE(entry->index);

    DEBUGK("UNLINK from 0x%p to 0x%p\n", vaddr, paddr);
    put_page(paddr);

    flush_tlb(vaddr);
}

// 拷贝一页，返回拷贝后的物理地址
static u32 copy_page(void *page)
{
    u32 paddr = get_page();

    page_entry_t *entry = get_pte(0, FALSE);
    entry_init(entry, IDX(paddr));
    memcpy((void *)0, (void *)page, PAGE_SIZE);

    entry->present = FALSE;
    return paddr;
}

typedef struct page_error_code_t
{
    u16 present : 1;
    u16 write : 1;
    u16 user : 1;
    u16 reserved0 : 1;
    u16 fetch : 1;
    u16 protection : 1;
    u16 shadow : 1;
    u16 reserved1 : 8;
    u16 sgx : 1;
    u16 reserved2;
} _packed page_error_code_t;

void page_fault(
        u32 vector,
        u32 edi, u32 esi, u32 ebp, u32 esp,
        u32 ebx, u32 edx, u32 ecx, u32 eax,
        u32 gs, u32 fs, u32 es, u32 ds,
        u32 vector0, u32 error, u32 eip, u32 cs, u32 eflags)
{
    assert(vector == 0xE);
    u32 vaddr = get_cr2();
    DEBUGK("fault address 0x%p\n", vaddr);

    page_error_code_t *code = (page_error_code_t *)&error;
    task_t *task = running_task();

    assert(KERNEL_MEMORY_SIZE <= vaddr && vaddr < USER_STACK_TOP);

    if (code->present)
    {
        assert(code->write);

        page_entry_t *pte = get_pte(vaddr, FALSE);
        page_entry_t *entry = &pte[TIDX(vaddr)];

        assert(entry->present);
        assert(memory_map[entry->index] > 0);
        if (memory_map[entry->index] == 1)
        {
            entry->write = TRUE;
            DEBUGK("WRITE page for 0x%p\n", vaddr);
        }
        else
        {
            void *page = (void *)PAGE(IDX(vaddr));
            u32 paddr = copy_page(page);
            memory_map[entry->index]--;
            entry_init(entry, IDX(paddr));
            flush_tlb(vaddr);
            DEBUGK("COPY page for 0x%p\n", vaddr);
        }
        return;
    }

    if (!code->present && (vaddr < task->brk || vaddr >= USER_STACK_BOTTOM))
    {
        u32 page = PAGE(IDX(vaddr));
        link_page(page);
        return;
    }
    panic("page fault!!!");
}

// 拷贝当前页目录
page_entry_t *copy_pde()
{
    task_t *task = running_task();

    page_entry_t *pde = (page_entry_t *)alloc_kpage(1);
    memcpy(pde, (void *)task->pde, PAGE_SIZE);

    // 将最后一个页表指向页目录自己，方便修改
    page_entry_t *entry = &pde[1023];
    entry_init(entry, IDX(pde));

    page_entry_t *dentry;

    for (size_t didx = (sizeof(KERNEL_PAGE_TABLE) / 4); didx < 1023; didx++)
    {
        dentry = &pde[didx];
        if (!dentry->present)
            continue;

        page_entry_t *pte = (page_entry_t *)(PDE_MASK | (didx << 12));

        for (size_t tidx = 0; tidx < 1024; tidx++)
        {
            entry = &pte[tidx];
            if (!entry->present)
                continue;

            // 对应物理内存引用大于 0
            assert(memory_map[entry->index] > 0);
            // 置为只读
            entry->write = FALSE;
            // 对应物理页引用加 1
            memory_map[entry->index]++;

            assert(memory_map[entry->index] < 255);
        }

        u32 paddr = copy_page(pte);
        dentry->index = IDX(paddr);
    }

    set_cr3(task->pde);

    return pde;
}

// 释放当前页目录
void free_pde()
{
    task_t *task = running_task();
    assert(task->uid != KERNEL_USER);

    page_entry_t *pde = get_pde();

    for (size_t didx = 2; didx < 1023; didx++)
    {
        page_entry_t *dentry = &pde[didx];
        if (!dentry->present)
        {
            continue;
        }

        page_entry_t *pte = (page_entry_t *)(PDE_MASK | (didx << 12));

        for (size_t tidx = 0; tidx < 1024; tidx++)
        {
            page_entry_t *entry = &pte[tidx];
            if (!entry->present)
            {
                continue;
            }

            assert(memory_map[entry->index] > 0);
            put_page(PAGE(entry->index));
        }

        // 释放页表
        put_page(PAGE(dentry->index));
    }

    // 释放页目录
    free_kpage(task->pde, 1);
    DEBUGK("free pages %d\n", free_pages);
}

int32 sys_brk(void *addr)
{
    DEBUGK("task brk 0x%p\n", addr);
    u32 brk = (u32)addr;
    ASSERT_PAGE(brk);

    task_t *task = running_task();
    assert(task->uid != KERNEL_USER);

    assert(KERNEL_MEMORY_SIZE <= brk && brk < USER_STACK_BOTTOM);

    u32 old_brk = task->brk;

    if (old_brk > brk) {
        for (; brk < old_brk; brk += PAGE_SIZE) {
            unlink_page(brk);
        }
    } else if (IDX(brk - old_brk) > free_pages) {
        // out of memory
        return -1;
    }

    task->brk = brk;
    return 0;
}